package com.tony.jdbc.dynamic.config;

import com.tony.jdbc.dynamic.datasource.DynamicDataSourceRouter;
import com.zaxxer.hikari.HikariDataSource;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.boot.autoconfigure.jdbc.DataSourceProperties;
import org.springframework.boot.jta.atomikos.AtomikosDataSourceBean;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Primary;
import org.springframework.jdbc.core.JdbcTemplate;

import javax.sql.DataSource;
import java.io.Closeable;
import java.io.IOException;
import java.sql.SQLException;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.stream.Collectors;

/**
 * @author tony
 * 普通多数据源配置、默认数据源和动态数据源合并，参考DataSourceAutoConfiguration
 */
@Configuration
@ConditionalOnProperty(prefix = "spring.jta", name = "enabled", havingValue = "false", matchIfMissing = true)
public class DynamicDataSourceAndTemplateConfig {
    private static final Logger LOGGER = LoggerFactory.getLogger(DynamicDataSourceAndTemplateConfig.class);
    private final DynamicDataSourceProperties dynamicDataSourceProperties;
    private JdbcTemplate jdbcTemplate;
    /**
     * 所有需要切换的数据源
     */
    private final Map<String, DataSource> dynamicDataSourceMap = new ConcurrentHashMap<>();

    /**
     * 构造函数
     *
     * @param dynamicDataSourceProperties 数据源配置信息
     */
    public DynamicDataSourceAndTemplateConfig(DynamicDataSourceProperties dynamicDataSourceProperties) {
        this.dynamicDataSourceProperties = dynamicDataSourceProperties;
    }

    @Primary
    @Bean
    public DynamicDataSourceRouter dynamicDataSource(DataSourceProperties dataSourceProperties) {
        DataSource defaultDatasource = dataSourceProperties.initializeDataSourceBuilder().build();
        jdbcTemplate = new JdbcTemplate(defaultDatasource);
        DynamicDataSourceRouter router = new DynamicDataSourceRouter();
        router.setDefaultTargetDataSource(defaultDatasource);
        this.updateAllDynamicDataSource(router);
        return router;
    }

    /**
     * 如果不存在，添加一个新数据源
     *
     * @param properties 数据连接信息
     * @return 数据源
     * @throws SQLException
     */
    private DataSource createIfNotExistsDataSource(DataSourceProperties properties) throws SQLException {
        String key = getKey(properties);
        DataSource dataSource = dynamicDataSourceMap.get(key);
        if (dataSource == null) {
            synchronized (dynamicDataSourceMap) {
                if (dynamicDataSourceMap.get(key) == null) {
                    //dataSource = DataSourceBuilder.derivedFrom(this).type(HikariDataSource.class).driverClassName(properties.getDriverClassName()).url(properties.getUrl()).username(properties.getUsername()).password(properties.getPassword()).build();
                    dataSource = createDataSource(properties, properties.getType());
                    dynamicDataSourceMap.put(key, dataSource);
                } else {
                    dataSource = dynamicDataSourceMap.get(key);
                }
            }
        }
        return dataSource;
    }

    @SuppressWarnings("unchecked")
    protected static <T> T createDataSource(DataSourceProperties properties, Class<? extends DataSource> type) {
        return (T) properties.initializeDataSourceBuilder().type(type).build();
    }

    /**
     * 更新数据源信息
     */
    public void updateAllDynamicDataSource(DynamicDataSourceRouter dynamicDataSourceRouter) {
        dynamicDataSourceRouter.setTargetDataSources(new HashMap<>());
        this.updateAllDynamicDataSource();
        dynamicDataSourceRouter.setTargetDataSources(new HashMap<>(dynamicDataSourceMap));
    }

    /**
     * 更新数据源信息
     */
    public void updateAllDynamicDataSource() {
        this.updateDynamicDataSource(dynamicDataSourceProperties.getAllDatasourceList(jdbcTemplate));
    }

    /**
     * 获取数据源key
     *
     * @param properties 数据源信息
     * @return 数据源key
     */
    private String getKey(DataSourceProperties properties) {
        return properties.getName();
    }

    /**
     * 更新数据源信息
     */
    public void updateDynamicDataSource(List<DataSourceProperties> newDatasourceList) {
        if (newDatasourceList == null) {
            //清理之前的数据源
            dynamicDataSourceMap.values().forEach(DynamicDataSourceAndTemplateConfig::closeDatasource);
            dynamicDataSourceMap.clear();
        } else {
            newDatasourceList.forEach(row -> {
                try {
                    dynamicDataSourceMap.putIfAbsent(getKey(row), createIfNotExistsDataSource(row));
                } catch (SQLException e) {
                    LOGGER.error("DynamicDataSource init datasource exception:{}", e.getMessage(), e);
                }
            });
            //通过key,删除未用到的。
            if (!dynamicDataSourceMap.isEmpty()) {
                Set<String> keySet = newDatasourceList.stream().map(this::getKey).collect(Collectors.toSet());
                Set<String> oldKeyIt = dynamicDataSourceMap.keySet();
                for (String key : oldKeyIt) {
                    if (!keySet.contains(key)) {
                        DataSource remove = dynamicDataSourceMap.remove(key);
                        if (remove != null) {
                            closeDatasource(remove);
                        }
                    }
                }
            }
        }
    }

    /**
     * 关闭数据源
     *
     * @param remove 数据源
     */
    private static void closeDatasource(DataSource remove) {
        if (remove instanceof AtomikosDataSourceBean){
            ((AtomikosDataSourceBean) remove).close();
        }else if (remove instanceof HikariDataSource){
            ((HikariDataSource) remove).close();
        }else if(remove instanceof Closeable){
            try {
                ((Closeable) remove).close();
            } catch (IOException e) {
                LOGGER.error("DynamicDataSource close datasource exception:{}", e.getMessage(), e);
            }
        }
    }

}
